import os
import pandas as pd
import numpy as np
from PIL import Image
import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

class Plain_Dataset(Dataset):
    def __init__(self,txt_file_path,img_dir,transform):
        '''
        Pytorch Dataset class
        params:-
                 txt_file_path : the path of the labels file    (train, validation, test)
                 img_dir       : the directory of the images (train, validation, test)
                 transform     : pytorch transformation over the data
        return :-
                 image, labels
        '''        
        labels = []
        # labels_oh_list = []
        # f = open(txt_file_path, 'r')
        # s = f.readlines()
        # f.close()
        with open(txt_file_path, 'r') as f:
            s = f.readlines()
        for i in range(len(s)):
            txt = s[i].split(' ')
            labels.append(np.array(txt[1][0], dtype='int8'))
        # labels_one_hot = pd.get_dummies(labels)
        # for idx, data in labels_one_hot.iterrows():
        #     label = np.array(data)
        #     labels_oh_list.append(label)
        # self.labels = labels_oh_list
        self.labels = labels
        self.img_dir = img_dir
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self,idx): 
        if torch.is_tensor(idx):
            idx = idx.tolist()
        img = Image.open(self.img_dir + str(idx)+'.jpg')
        labels = self.labels[idx]
        labels = torch.from_numpy(labels).long()
        if self.transform :
            img = self.transform(img)
        return img,labels

#Helper function
def eval_data_dataloader(txt_file_path,img_dir,datatype,sample_number,transform= None):
    '''
    Helper function used to evaluate the Dataset class
    params:-
            txt_file_path : the path of the labels file    (train, validation, test)
            img_dir       : the directory of the images (train, validation, test)
            sample_number : any number from the data to be shown
    '''
    if transform is None :
        transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,),(0.5,))])
    dataset = Plain_Dataset(txt_file=txt_file,img_dir = img_dir, transform = transform)

    label = dataset.__getitem__(sample_number)[1]
    print(label)
    imgg = dataset.__getitem__(sample_number)[0]
    imgnumpy = imgg.numpy()
    imgt = imgnumpy.squeeze()
    plt.imshow(imgt)
    plt.show()
